/*
 * Huffman encoding test generator.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#define TEST_MAXLEN	100000

#define SMLTESTS	15
#define LRGTESTS	50

char *
make_random_alphabet(int alpha)
{
	char *a;
	int i = 0;
	int j;

	if (alpha > 26 || alpha < 1) {
		fprintf(stderr, "%s: bad argument (alpha=%d)\n", __func__, alpha);
		return NULL;
	}

	a = calloc(alpha, sizeof *a);
	while (i < alpha) {
		a[i] = 'a' + rand() % 26;
		j = 0;
		while (j < i) {
			if (a[j] == a[i])
				j = 0, a[i] = 'a' + rand() % 26;
			else
				j++;
		}
		i++;
	}

	return a;
}

char *
freq_string(int alpha, char *alphabet, int *freq, int len)
{
	int *freq_rev;
	int s, i, c;
	char *str;

	for (s = i = 0; i < alpha; i++)
		s += freq[i];
	str = calloc(s + 1UL, sizeof *str);
	freq_rev = calloc(alpha, sizeof *freq_rev);

	c = 0;
	while (c < len) {
		i = s = rand() % alpha;
		while (freq_rev[i % alpha] == freq[i % alpha] &&
				i < s + alpha)
			i++;
		if (i == s + alpha)
			break; /* done */
		i %= alpha;
		freq_rev[i]++;
		str[c++] = alphabet[i];
	}
	if (c == len)
		fprintf(stderr, "%s: string ended early\n", __func__);
	str[c++] = '\0';

	free(freq_rev);
	return str;
}

char *
chained_tree_test(int seed, int alpha, int startsize, int step, int size)
{
	int i;
	char *alphabet;
	int freq[26];
	char *str;

	if (alpha < 1 || alpha > 26 || startsize < 0 || step < 0 || size < 1) {
		fprintf(stderr, "%s: invalid arguments\n", __func__);
		return NULL;
	}

	srand(seed);
	memset(freq, 0, sizeof freq);
	alphabet = make_random_alphabet(alpha);
	freq[0] = startsize;
	freq[1] = startsize + step;
	for (i = 2; i < alpha; i++)
		freq[i] = freq[i - 1] + freq[i - 2];
	if (freq[alpha - 1] + freq[alpha - 2] + freq[alpha - 1] >= TEST_MAXLEN) {
		fprintf(stderr, "%s: test length estimation exceeded TEST_MAXLEN (%d)\n",
				__func__, TEST_MAXLEN);
	}
	str = freq_string(alpha, alphabet, freq, size);
	free(alphabet);

	return str;
}

char *
balanced_tree_test(int seed, int alpha, int size)
{
	int i;
	char *str;

	srand(seed);
	str = malloc(size + 1UL);
	for (i = 0; i < size; i++)
		str[i] = 'a' + (rand() % alpha);
	str[alpha] = '\0';

	return str;
}

char *
generate_test(int seed, int alpha, int size)
{
	int freq[26];
	char *alphabet;
	int i;
	char *str;

	srand(seed);
	alphabet = make_random_alphabet(alpha);
	memset(freq, 0, sizeof freq);
	for (i = 0; i < alpha; i++)
		freq[i] = 1 + (rand() % (TEST_MAXLEN / alpha));

	str = freq_string(alpha, alphabet, freq, size);
	free(alphabet);
	return str;
}

int
main(int argc, char **argv)
{
	char *tests[SMLTESTS + LRGTESTS];
	char filename[128];
	int i;
	FILE *fp;

	tests[0] = strdup("stress");
	tests[1] = strdup("a");
	tests[2] = strdup("aaaaaaa");
	tests[3] = strdup("bbbbbbbb");
	tests[4] = strdup("ccccccccc");
	tests[5] = strdup("ddddddde");
	tests[6] = strdup("dddddddde");
	tests[7] = strdup("abbabaab");
	tests[8] = chained_tree_test(6, 10, 1, 0, 300);
	for (i = 9; i < SMLTESTS + LRGTESTS - 5; i++)
		if ((i + 1) % 8 == 0 && (i + 1) % 3 == 0)
			tests[i] = balanced_tree_test(i, 2 + (rand() % 24), 300 + (rand() % (TEST_MAXLEN - 301)));
		else if ((i + 1) % 8 == 0)
			tests[i] = chained_tree_test(i, (rand() % 10) + 16, rand() % 3, rand() % 2, TEST_MAXLEN);
		else
			tests[i] = generate_test(i, 1 + (rand() % 25), 1 + (rand() % TEST_MAXLEN));
	tests[SMLTESTS + LRGTESTS - 5] = chained_tree_test(i, 26, 1, 0, TEST_MAXLEN);
	tests[SMLTESTS + LRGTESTS - 4] = generate_test(i, 1, TEST_MAXLEN);
	tests[SMLTESTS + LRGTESTS - 3] = generate_test(i, 2, TEST_MAXLEN);
	tests[SMLTESTS + LRGTESTS - 2] = generate_test(i, 3, TEST_MAXLEN);
	tests[SMLTESTS + LRGTESTS - 1] = generate_test(i, 26, TEST_MAXLEN);

	for (i = 0; i < SMLTESTS + LRGTESTS; i++) {
		snprintf(filename, sizeof filename, "%02d.txt", i + 1);
		if (fp = fopen(filename, "w+t"), fp == NULL) {
			fprintf(stderr, "%s: Failed to open file \"%s\" for R/W, skipping..\n",
					argv[0], filename);
		} else {
			fprintf(fp, "%s\n", tests[i]);
			fflush(fp);
			fclose(fp);
			fp = NULL;
		}
	}

	for (i = 0; i < SMLTESTS + LRGTESTS; i++)
		free(tests[i]);

	return 0;
}
